Numerical model of quantum wavepacket dynamics
using the time-dependent Scrödinger equation:
Animation of tunneling through a barrier in 1D and 2D.
using the time-dependent Scrödinger equation:
Animation of tunneling through a barrier in 1D and 2D.
The program below models a particle wavefunction dynamics in the field of an arbitrary potential. In particular, tunneling through a square barrier of finite height and width is demonstrated, as well as particle evolution in a harmonic potential. The method is based on solving numerically the time-dependent Schrödinger equation (TISE). The results are visualized via animated plots.
We will present a brief theoretical background and numerical method to obtain the solutions of the time-dependent Schrödinger equation by using matrix/vector representation of operators and wavefunctions.
1. Theoretical basis¶
We will consider two stable methods for solving the TDSE recursively.
1.1. Solving TDSE in coordinate space¶
We start from the time-dependent Schrödinger equation: $$ \hat{H} \psi(x,t) = i \hbar \frac{\partial \psi(x,t)}{\partial t} $$ Using the definition of the derivative, and swapping the sides of the equation, we rewrite it as: $$ \frac{\psi(x,t+dt)-\psi(x,t)}{dt} = \frac{1}{i \hbar}\hat H\psi(x,t) $$ An explicit method (or forward Euler method) is to recursively calculate the wavefunction at a later time $\psi(t+\Delta t)$ from the wavefunction at an earlier time $\psi(t)$ using the equation above with small increments of time $\Delta t$: $$ \psi(x,t+\Delta t) \approx \psi(x,t) + \frac{1}{i \hbar}\left [\hat{H} \psi(x,t) \right ] \Delta t $$
Unfortunately, the explicit method shown above is unstable; small errors in numerical calculations accumulate in each recursive time step, and the solution diverges.
Several workarounds have been proposed.
One of them involves combining the explicit forward approach (shown above) with the backward Euler (implicit) method, where the value of the desired wavefunction at the next time step is expressed from its value at that step (Physics, Python, and Programming).
$$
\frac {\psi(x,t+dt)-\psi(x,t)}{dt} \approx \frac{1}{i\hbar}\hat H\psi(x,t+dt)
$$
Averaging the explicit and implicit methods leads to the Crank-Nicolson scheme, which combines adequate accuracy and stability at a modest computational cost:
$$
\frac {\psi(x,t+dt)-\psi(x,t)}{dt} \approx \frac{1}{i \hbar}\frac{\hat H\psi(x,t) + \hat H\psi(x,t+dt)}{2}
$$
In the equation above, the rate of change of the function $\psi$ from $t$ to $t+\Delta t$ is just an average of the derivatives at time $t$ and $t+\Delta t$, which approximates the derivative between these two time points, hence the better precision and stability of the method.
Expressing $\psi(x,t+dt)$ from the above equation results in the final equation:
$$
\psi(x,t+dt) = \left [1-\frac{dt}{2i\hbar}\hat H \right ]^{-1} \left [1+\frac{dt}{2i\hbar}\hat H \right ]\psi(x,t)
$$
1.2. Solving TDSE in coordinate and momentum space¶
The alternative approach is to use the potential evolution in real space, and use the kinetic evolution in Fourier space
(Sushata Barman. Quantum Tunneling: Theoretical Insights and Numerical Simulation).
We start from the Schrödinger equation:
$$
\hat{H} \psi(x,t) = i \hbar \frac{\partial \psi(x,t)}{\partial t}
$$
or:
$$
\frac{\partial \psi(x,t)}{\partial t} = \left ( -\frac{i\hat H}{\hbar} \right ) \psi(x,t)
$$
We will assume further that during a short enough time $\Delta t$, the action of the Hamiltonian $\hat H = \hat T + V$ on the wavefunction is constant. In that case, the equation states that the time derivative of a function is proportional to the function itself times a constant, and the solution is an exponential function:
$$
\psi(x, t + \Delta t) \approx \exp \left ( -\frac{i\hat H}{\hbar} \Delta t\right ) \psi(x,t) \\
\psi(x, t + \Delta t) \approx \exp \left ( -\frac{i\hat T}{\hbar} \Delta t\right ) \exp \left ( -\frac{iV}{\hbar} \Delta t\right )\psi(x,t)
$$
The equation above allows us to calculate $\psi(x,t)$ recursively using small increments of time.
Note that the actions of kinetic $\hat T$ and potential $V$ energies are multiplicative and they can be performed in two separate steps. The stable recursive method involves three steps in a single iteration propagating the wavefunction for time $\Delta t$. The action of the potential energy is performed in the x-coordinate basis, and the action of the kinetic energy is performed in the momentum space, where the kinetic energy is simply $\hat T = \hat p^2/2m.$
STEP 1.
In the first step, the wavefunction is propagated in time for a half-step ($\Delta t/2$) using the factor involving the potential energy V:
$$
\psi (x_i,t+ \Delta t/2) = \exp\left ( -\frac {iV(x)\Delta t}{2\hbar}\right) \psi (x_i,t) \qquad (1)\\
$$
Note that the potential energy can be expressed numerically as a matrix with its elements being V(x), and this step is thus easily realized numerically.
STEP 2.
In the second step, the halfway function obtained in the first step is propagated in time for a full step ($\Delta t$) using the kinetic energy term. In that step, the wavefunction is expressed in the space of the eigenfunctions of the momentum operator, which are the solutions to the equation:
$$
\hat p \phi(x) = p\phi(x) \\
-i\hbar \frac{d\phi(x)}{dx} = p \phi(x) \\
\frac{\phi(x)}{dx} = \frac {ip}{\hbar} \phi(x)
$$
i.e., the derivative of the function is proportional to the function itself. The solutions to this equation are, naturally, plane waves:
$$
\phi(x) = A\exp(ipx/\hbar)
$$
We will omit the scaling (normalization) constant $A$ for brevity, assuming that care is taken to normalize the wavefunctions before use.
If x-space is confined to length $L$, the eigenvalues $p_n$ are quantized due to the boundary conditions at the ends of the x-space (a particle in a box), i.e., the number of half-waves between the points must be an integer:
$$
p_n = \hbar k = \frac{2 \pi n \hbar}{L} \\
\phi_n(x) = \exp(i p_n x/\hbar)
$$
Note that in the case of a discrete x-space, the number of eigenfunctions is finite (equal to the number of grid points in x).
To transform the wavefunction from x-space to p-space, one has to expand the wavefunction $\psi(x)$ in the basis of the momentum operator eigenfunctions using the Fourier trick:
$$
\psi(x) = \sum_n {b_n \phi_n(x)} \\
b_n = \left < \psi| \phi_n \right> = \int {\psi^*(x) \exp(i p_n x/\hbar) dx }
$$
The latter is identical to the Fourier transform of the function $\psi(x)$ to $\psi(p)$, i.e., the wavefunction vector $\psi(x_i)$ is transformed into $\psi(p_i)$ in p-coordinate system. Respectively, the wavefunction transformation of the wavefunction from p-space back to x-space involves the inverse Fourier transform.
The momentum operator in the basis of its eigenfunctions is just a diagonal matrix of eigenvalues of p (i.e., $p_n$), and thus the kinetic energy is easily found as $\hat T = \hat p^2/2m = (\hbar k)^2/2m$. We now propagate the half-step wavefunction to a full $\Delta t$ step in p-space using the kinetic energy:
$$
\widetilde \psi ' (k,t+ \Delta t) = \exp \left [ -\frac{i \hbar k^2 \Delta t}{2m}\right ] \widetilde \psi (k,t+\Delta t/2) \qquad (2)
$$
STEP 3.
in the final step, the wavefunction is transformed back into x-space and propagated in time for an additional half-step using the potential energy:
$$
\psi (x_i,t+ \Delta t) = \exp\left [ -\frac {iV(x)\Delta t}{2\hbar}\right] \psi ' (x_i,t+ \Delta t) \qquad (3)
$$
Here, the $\widetilde \psi(k_x,t)$ is the Fourier transform of the $\psi(x,t)$.
As a result of these three steps, we propagate the wavefunction in time for a full $\Delta t$ step using both kinetic and potential energy terms. Dividing time propagation into three steps ensures that the solution does not diverge after multiple recursive $\Delta t$ increments.
2. Numerical solution of the Schrödinger equation¶
2.1. Vector representation of a wavefunction.¶
In numerical modeling, the wavefunction is represented as a discrete function: $$\psi(x,t) = \psi(x_i ,t) \equiv \psi _i$$ where x is an array of x-positions $\{x_i\}$ at which the function array $\{\psi _i\}$ is defined or calculated. Assuming $\{x_i\}$ is a vector of size N: $$\mathbf x = \left [ x_1, x_2, ..., x_N \right ] $$ The wavefunction is also a vector whose size is defined by the number of points in the array $\{x_i\}$: $$ \boldsymbol \psi(x,t) = \left [ \psi(x_1 ,t), \psi(x_2 ,t), ..., \psi(x_N ,t) \right ] \equiv \left [\psi_1, \psi_2, ..., \psi_N \right ] $$ Note that care should be taken in numerical multiplications of that vector with other vectors or matrices; in some cases, the vector should be a row, in others it should be a column.
2.2. Matrix representation of the potential energy operator¶
The time-independent Schrödinger equation involves the product of the potential energy with the wavefunction:
$$
\hat{H} \varphi(x) = \frac{\hbar}{2m}\frac{d^{2}\varphi(x)}{dx^{2}} + V(x)\varphi(x)
$$
Since we work in a discrete space of $\{ x_i \}$, the product $V(x)\varphi(x) = V(x_i)\varphi(x_i) \equiv V_i\varphi_i$.
The potential function is, thus, also represented as a set of values $\{ V_i \}$ at positions $\{x_i\}$. For computational efficiency, we will represent the potential energy operator as a diagonal matrix $\mathbf V$ with its diagonal values being $\mathbf V_{ii} \equiv V(x_i) $. It can be easily shown that the product $V_i\varphi_i$ can then be represented as the following matrix product:
$$
V(x) \psi (x) = \mathbf{V}\boldsymbol \varphi = \begin{pmatrix}
&V(x_1) &0 &0 &... \\
&0 &V(x_2) &0 &... \\
&0 &0 &V(x_3) &... \\
&... &... &... &... \\
\end{pmatrix}
\begin{pmatrix}
\varphi(x_1) \\
\varphi(x_2) \\
\varphi(x_3) \\
... \\
\end{pmatrix}
= \begin{pmatrix}
V_1\varphi_1 \\
V_2\varphi_2 \\
V_3\varphi_3 \\
... \\
\end{pmatrix}
$$
The above matrix product multiplies the wavefunction vector by the potential at each point - the goal is achieved.
The right side of the time-independent Schrödinger equation $E \psi (x)$ involves the product of a constant by the wavefunction vector, and can also be represented as a matrix product: $$ E \psi (x) = E\mathbf I \boldsymbol \varphi = \begin{pmatrix} &E &0 &0 &... \\ &0 &E &0 &... \\ &0 &0 &E &... \\ &... &... &... &... \\ \end{pmatrix} \begin{pmatrix} \varphi(x_1) \\ \varphi(x_2) \\ \varphi(x_3) \\ ... \\ \end{pmatrix} = \begin{pmatrix} E \varphi_1 \\ E \varphi_2 \\ E \varphi_3 \\ ... \\ \end{pmatrix} $$ where $\mathbf I$ is the identity matrix (unit matrix), i.e., the diagonal matrix with its diagonal elements being equal to 1. The above expression multiplies all components of vector $\boldsymbol \varphi$ by a scalar $E$.
2.3. Kinetic energy matrix in the momentum space¶
In the momentum space, the kinetic energy operator is simply a square of the momentum, $\hat T = \hat p^2/2m = (\hbar k)^2/2m$. The transformation of a wavefunction from x-space to p-space involves the Fourier transform, which is a standard function in the numpy library (numpy.fft.fft). Note that fft stands for Fast Fourier Transform, and requires that the number of points in the function is $2^m$, where $m$ is an integer number. In other words, the number of grid points in the x-axis must be a multiple of 2, such as 128, 256, etc.
The eigenvalues of $\hat p$ are defined by the discrete frequencies $f_n$ of the Fourier transform, which can be easily obtained using the function numpy.fft.fftfreq, and $p_n$ eigenvalues are:
$$
p_n = \hbar k = \hbar 2 \pi f_n
$$
In the matrix form, $\mathbf p$ is a diagonal matrix with the eigenvalues $p_n$ on its diagonal,
and the kinetic energy is square of this matrix:
$$
\mathbf T = \mathbf p^2 / 2m
$$
The Fourier inverse transform is performed by the function numpy.fft.ifft.
2.4. Modeling wavefunction dynamics using TDSE equation¶
Once the operators and wavefunction are represented numerically as matrices and vectors, the TDSE solution will involve matrix algebra. In the first method, numerical calculations in matrix form are given by: $$ \boldsymbol \psi(t+\Delta t) = \boldsymbol \psi(t+\Delta t) = \left [1-\frac{dt}{2i\hbar}\mathbf H \right ]^{-1} \left [1+\frac{dt}{2i\hbar}\mathbf H \right ]\boldsymbol \psi(t) $$ Similarly, the second method also involves matrix operations: $$ \boldsymbol \psi (x_i,t+ \Delta t/2) = \exp\left ( -\frac {i\Delta t}{2\hbar}\mathbf V \right) \boldsymbol \psi (x_i,t) \qquad (1)\\ $$ $$ \boldsymbol {\widetilde \psi} ' (k_i,t+ \Delta t) = \exp \left [ -\frac{i \hbar \Delta t}{2m} \mathbf k^2 \right ] \widetilde {\boldsymbol \psi} (k_i,t+\Delta t/2) \qquad (2) $$
$$ \boldsymbol \psi (x_i,t+ \Delta t) = \exp\left [ -\frac {i\Delta t}{2\hbar} \mathbf V\right] \boldsymbol \psi ' (x_i,t+ \Delta t) \qquad (3) $$All these operations can be performed using standard functions found in the numpy library of Python.
3.1.1. Import the necessary libraries and set some parameters.¶
The import statement tells Python to load a library of functions. The statement can be followed by an optional as keyword that allows one to substitute the name of the original library to simpilfy its use.
The from statement allows importing a specific function from a library; once imported this functions can be called by its name without providing a full path. The * will import all functions from the library without the need to use the full path.
Finally, we set $\hbar =1$ and $m=1$, i.e., we will work in atomic units and assume that the mass of the particle is that of an electron.
*save_mp4*:
set it to False if you do not have ffmpeg installed separately and its bin added to system's PATH. ffmpeg is used to create mp4 animation file, the program will fail if that is not the case.
*show_html*:
set it to False if you only want to get mp4 movie. If True, shows the animation in a browser, assuming that you run under Jupyter. May fail in stand-alone Python that does not output to browser (never tested).
*use_latex*:
When True, the animation will have equations written on them, but that require an external LaTeX package. The program will fail if it is not the case, so set that to False. You can install a LaTeX package separately, such as MiKTeX, and provide system PATH to its bin directory
# By Sergei Savikhin, Purdue University, September 2025
import numpy as np # 'numerical python' library, defines matrices, comlpex numbers etc, we will refer to its functions using prefix 'np'
import matplotlib.pyplot as plt # set of functions for drawing graphs
from matplotlib.animation import FuncAnimation # set of functions to create animations using set of plots
from IPython.display import HTML, display # to export instructions in html format (inline Jupiter output)
from math import * # input all math functions (sin, cos, exp...), with no need to use any prefix
hbar = 1 # set hbar to 1
m = 1 # set mass to 1
save_mp4 = True # Save animation as MP4 file. Requires ffmpeg codec to be installed separately
# and its bin directory must be added to the system PATH to be found by Python
show_html = True # embed animation into output brtowser window. Will work under Jupyter
# but probably not in a stand alone Python like Spyder
3.1.3. Build the initial Gaussian wavepacket:¶
We will generate the initial wavefunction $\psi(x,t=0)$ as a complex Gaussian wavepacket: $$ \psi(x,t=0) = A \exp{\left [-\frac{(x - x_0 )^2}{2 \sigma ^2}\right ]} \exp{(ik_0x)} $$ Here, the first term is the Gaussian function that defines the amplitude of the wave; its width is defined by $\sigma$, and the maximum position is at $x_0$. The second term is a complex oscillating function with wavenumber $k_0$. The wavefunction will be normalized by picking a correct A so that the integral of the $|\psi|^2$ over the whole space of x gives exactly 1: $$ \left < \psi(x) | \psi(x) \right > = \int {\left | \psi \right |^2dx} \approx \sum_i {\left | \psi_i \right |^2 \Delta x} = 1 $$ There is a connection between the k-number and the kinetic energy of the particle wave, $E_{wave} = \frac {p^2}{2m} = \frac {(\hbar k)^2}{2m}$. Therefore, a particle's wavepacket can be defined by either center wavevector $k_0$ as in the equation above, or by its kinetic energy $E_{particle}$. The expectation energy of the particle, $\left<E\right>$, described by the wavepacket that is symmetric along x, however, is not exactly equal to $(\hbar k)^2 /2m$ since the energy increases with k. One can obtain $k_0$ for the particle's expectation energy $\left<E\right>$ as:
$$ k_0 = \frac {m}{\hbar}\sqrt {2 \left < E \right >^2 - \frac {1}{2 \sigma^2}} $$Note that the kinetic energy of the particle could be found using: $$ \left < T \right > = \left < \psi | T | \psi \right > = \boldsymbol \psi^T \mathbf T \boldsymbol \psi $$ Since matrix calculations run in discrete space, the energy calculated using the above equation and respective matrices may yield slightly different results compared to the expectation energy used to set $k_0$.
Input parameters:
x = is an array of x-coordinates $\{x_i\}$ where the wavefunction is calculated. If no parameter is provided, it is set to range from 0 to 1000 and step $\Delta x$ =1
x0 = the position of the center of the wavepacket along x, i.e. $x_0$, if not provided the default value 200 is used
k0 - wavenumber of the central wave component (default 0.4)
sigma = $\sigma$, the width of the Gaussian (default is 15)
Note: for the second function, one defines the expectation energy of the particle E instead of k0
Output:
$\psi(x,t=0)$ array, $\{\psi_i\}$
Notes:
- np.arange(x1,x2,dx) - numpy function, returns an array of values ranging from x1 to x2-dx, with step dx
- np.exp(x) - returns an exponent of an array (new array = $\exp(x_i)$)
- 1j $\equiv \sqrt{-1} \equiv i$
- numpy functions are 'smart' and the exact action depends on the arguments. For example, in the expression (x - k0), x is an array, while k0 is a scalar. The result will be an array where the scalar k0 is subtracted from each component of x. Also, when np.exp is used on a complex number, the result is also a complex number (or array).
def Gaussian_wave(x = np.arange(0,1000), x0 = 200, k0 = 0.4,sigma = 15):
wf = np.exp( -1/2* (x-x0)**2/sigma**2) *np.exp(1j*k0*x)
wf = wf/(sqrt(np.sum(np.abs(wf)**2)*(x[1]-x[0]))) # normalize
return wf
def Gaussian_wave_at_E(x = np.arange(0,1000), x0 = 200, Ep = 0.1,sigma = 15):
k0 = m/hbar * sqrt(2*(Ep-1/sigma**2/4))
return Gaussian_wave(x = x, x0 = x0, k0 = k0, sigma = sigma)
3.1.4. Build the Potential energy matrix¶
As described in the Theory section, the potential energy operator is represented by a diagonal matrix :
$$
\mathbf V = \begin{pmatrix}
&V(x_1) &0 &0 &... \\
&0 &V(x_2) &0 &... \\
&0 &0 &V(x_3) &... \\
&... &... &... &... \\
\end{pmatrix}
$$
Input: parameters:
x - the x range to calculate V
position - the position where the barrier starts
height - the energy on the top of the barrier
width - the width of the barrier along x
Returns a list of output objects
V - The potential energy matrix
flatV - the 1D array of $V(x)$, where $V=0$ everywhere except between position and position+width, where its value V =height.
def Potential_barrier(x,position,height,width):
flatV = np.array([height if position< pos < position+width else (0) for pos in x]) #make 1D array
V = np.diag(flatV) #make matrix with given diagonal
return V,flatV
3.1.5. Build the kinetic energy matrix
The kinetic energy matrix is built according to the Finite Difference Method desribed in the Theory section.
$$
\mathbf T = -\frac{\hbar}{2m} \frac {1}{\Delta x^2}\begin{pmatrix}
&-2 &1 &0 &0 &0 &... &(1) \\
&1 &-2 &1 &0 &0 &... &0\\
&0 &1 &-2 &1 &0 &... &0\\
&0 &0 &1 &-2 &1 &... &0\\
&... &... &... &... &... &... &... \\
&(1) &0 &0 &0 &0 &1 &-2 \\
\end{pmatrix}
$$
The (1) in the opposite corner makes T periodic, as described earlier, the x-space will be periodic, when a free particle hits the edge of the $\{x_i\}$ range, it appears on the other side.
Input parameter:
x - the array of x-values corresponding to the range where the wavefunction is calculated. It is only used to find $\Delta x = x_1 - x_0 $ and the number of points (size) of the requested matrix
periodic - if True, the (1) will be added to the matrix for periodicity
Returns: The kinetic energy matrix
def Kinetic_energy(x,periodic = False):
size = len(x) # returns the size (number of points) in x-array
dx = x[1]-x[0] # difference between two nearby points
D2 = (np.diag(-2*np.ones(size)) #make a matrix with its main diagonal having size "size" and fill the diagonal with "-2" values
+ np.diag(np.ones(size-1),1) # make matrix with its diagonal one below the main diagonal, its size will be "size-1", and fill with '1's
+ np.diag(np.ones(size-1),-1)) # diagonal one above the main.
if periodic:
D2[size-1,0] = 1 # these two lines make the T (and Hamiltonian) periodic
D2[0,size-1] = 1 # in temrs of the derivative once it needs a point to the left of 0-point, it assumes it is the last point in the wavefunction
T = (-hbar/m/2) * (1/dx**2) * D2
return T
3.1.5. Test wavepacket and potential barrier functions (plot them)¶
We will generate a wavepacket and potential, and plot both real and imaginary parts of the wavefunction, the square of the wavefunction, and the potential barrier.
x = np.linspace(400,600,512) # make an x coordinate grid spanning from 400 to 600 with 512 points in between
wf = Gaussian_wave(x,x0 = 500) # generate in x grid, set max position to 500, and use default k0 and sigma
V,flatV = Potential_barrier(x,position=550,height=0.2,width=20) # geberate potential barrier
fig = plt.figure() # create figure
ax = plt.axes() # add axes to plot graphs
ax.plot(x,wf.real, color = 'red') # plot real part of wavefuncton s a function of x
ax.plot(x,wf.imag,color = 'blue') # plot imaginary part
ax.plot(x,np.square(np.abs(wf)),color = 'black') # plot square of the wavefunction (probability)
ax.plot(x,flatV,color = 'green') # plot potential barrier
ax.grid() # add grid to plot
plt.show() # show plot
plt.close() # close plot
3.1.6. Propagating wavefunction in time using method 1¶
To calculate the time evolution, the subroutine uses the following equation $$ \psi(x,t+dt) = \left [1-\frac{dt}{2i\hbar}\hat H \right ]^{-1} \left [1+\frac{dt}{2i\hbar}\hat H \right ]\psi(x,t) = \hat{M}_{evolution} \psi(x,t) \\ \hat{M}_{evolution} \equiv \left [1-\frac{dt}{2i\hbar}\hat H \right ]^{-1} \left [1+\frac{dt}{2i\hbar}\hat H \right ] $$ The subroutine will propagate the wavefunction from the initial time t0 recursively to the final time tf using time step dt; the returned function may be not be at the exact time tf since the time step is discrete.
Input:
psi0 - the initial wavefunction at time t0
t0 - time at which psi0 was recorded
evolutionmatrix = $M{evolution}$ matrix
tf - time at which wf needs to be calculated
Returns:
psi - new wavefunction
t0 - current time at which the function was calculated
def GetPsi2(psi0,t0,tf,evolution_matrix,dt):
psi = psi0 # initial psi at tme t0
#t = t0 # current time
while fabs(tf-t0) > dt/2 and t0 < tf: # continue recursively incrementin time by dt untlil time teaches tf
t0 += dt # new time point
psi = evolution_matrix @ psi # since we defined matrix and vector as numpy arrays ise @ sign to denote matrix dot-product
return psi, t0 # return wavefunction at new time and actual time
3.1.7. Wavepacket tunneling animation using method 1¶
x = np.linspace(0,800,800) # x-grid to work in
Ep, x0, sigma = 0.1, 300, 15 # the initial particle's wavefubction parameters
barrier_pos, barrier_height, barrier_width = 400, 0.2, 3 # barrier parameters
dt = 1 # dt value for simulation, the smaller will increase both precision and computational time
snapshot_times = [0,300] # calculate static snapshots at these times (set it to empty [] to skip)
t = np.arange(0,700,10) # animation frames calculated at these time values
T = Kinetic_energy(x)
V, flatV = Potential_barrier(x,barrier_pos,barrier_height,barrier_width)
H = T + V
wf0 = Gaussian_wave_at_E(x,x0 = 300,Ep = Ep,sigma=sigma)
I = np.identity(len(H))
evolution_matrix = np.linalg.inv(I - dt / 2.0j * H) @ (I + dt / 2.0j * H)
wf, curtime = wf0.copy(), 0 # pass a copy of wavefunction to wf to keep wf0 intact
if len(snapshot_times) > 0: # plot some snapshots of |psi|^2 if some snapshot_times are defined
print(f'working on snapshots...{snapshot_times}',end = '')
ax = plt.axes()
for tim in snapshot_times:
curtime = 0
wf,curtime = GetPsi2(wf,curtime,tim,evolution_matrix,dt)
ax.plot(x,np.square(np.abs(wf)))
plt.show()
plt.close()
wf, curtime = wf0.copy(), 0 # restore the original initial wavefunction
print('done')
# Create plot space
fig = plt.figure(figsize = (10,4) ,dpi=100) # make plot frame
ax = plt.axes(xlim = (np.min(x), np.max(x)),ylim = (-.3,.4))
ax.plot(x,flatV) # plot potential
ax.grid()
plt.title(f'Particle tunneling: Ep = {Ep}, Barrier height = {barrier_height} width = {barrier_width}')
line_wfreal, = ax.plot([],[],color = 'red') # create an empty linein plot, we can replace it with new line in the loop for animation
line_wfimag, = ax.plot([],[],color = 'blue') #placeholder for imaginary part of wf
line_wf2, = ax.plot([],[],color = 'black') # placeholfer for suqare of the wavefunction
last_frame = -1
wf, curtime = wf0.copy(), 0
def animate(frame): # this function will be executed by FuncAnimation to plot a new plot frame, i will refpresent framenumber to plot (i.e. 0,1,2,3...)
print(f'\rFrame {frame+1} of {len(t)}...',end='')
global wf,curtime # these are defined outside of the function and protected from overwriting. Making them global allows to change their values.
wf, curtime = GetPsi2(wf,curtime,t[frame],evolution_matrix,dt) # get new wavefunction at next time
line_wfreal.set_data(x, wf.real + Ep)
line_wfimag.set_data(x, wf.imag + Ep)
line_wf2.set_data(x, np.square(np.abs(wf))*5-.3)
ani = FuncAnimation(fig, animate, frames=len(t),interval=50)
plt.close()
if save_mp4 == True:
name = f'packet_E{Ep}_barrier_h{barrier_height}-w{barrier_width}-TDSE.mp4'
print(f'Saving in "{name}"...')
ani.save(name) # will save a movie in mp4 video file. Need to install ffmpeg codec and add its /bin/ to windows PATH!
print('done')
if show_html == True:
print('Making HTML video...')
wf, curtime = wf0.copy(), 0 #NB! ani.save actually plots and runs simulation, so we need to set initial conditions right!
anim = ani.to_jshtml()
print('done')
display(HTML(anim)) # for some reason, plain HTML(anim) does not work when within if block??? Crazy!
3.1.8. Propagate wavefunction in x- and p-space (GetPsi4 function)¶
We will first define the wavefunction propagation function GetPsi4 from time t0 to time tf using time increments dt and the three-step procedure described in the theory section:
Step 1: Apply half-step $\mathbf{potential}$ operator:
$$
\boldsymbol \psi (x_i,t+ \Delta t/2) = \mathbf{potential} \cdot \boldsymbol \psi (x_i,t) \qquad (1)\\
\mathbf {potential} \equiv \exp\left ( -\frac {i\Delta t}{2\hbar}\mathbf V \right)
$$
Step 2: Fourier trasform wavefunction into p-space, apply full-step kinetic operator in p-space, and fourier inverse wavefunction into x-space
$$
\boldsymbol {\widetilde \psi} ' (k_i,t+ \Delta t) = \mathbf{kinetic\_opertor} \cdot \widetilde {\boldsymbol \psi} (k_i,t+\Delta t/2) \qquad (2) \\
$$
$$
\mathbf {kinetic\_opertor} \equiv \exp \left [ -\frac{i \hbar \Delta t}{2m} \mathbf k^2\right ]
$$
Step 3: Apply half-step $\mathbf{potential}$ operator
from scipy.fftpack import fft, ifft, fftshift, ifftshift # will use fast fourier from science python pack
def GetPsi4(psi0,t0,tf,potential,kinetic_operator,dt):
psi = psi0 # initial psi at tme t0. Note: this does not make a copy of psi0, but merely assigns name psi to the same array
while fabs(tf-t0) > dt/2 and t0 < tf: # continue recursively incrementing time by dt until time reaches tf or we are within dt/2 of target time
#print(t0)
t0 += dt # new time point, increment by dt
# STEP 1: Apply potential operator (half step)
psi *= potential
# STEP 2: Apply kinetic operator (full step in Fourier space)
psi_k = fft(psi) # fourier transform wavefunction into p-space (k-space)
psi_k *= kinetic_operator # propagate by dt in momentum space
psi = ifft(psi_k) # fourier inverse wavefunction into x-space
# STEP 3: Apply potential operator (half step)
psi *= potential
return psi, t0 # return wavefunction at new time and new time itself
3.1.9. Tunneling animation using GetPsi4 function (x- and p- space)¶
# Simulation parameters
Lx = 800 # Spatial dimension of the coordinate space
Nx = 1024 # Number of grid points along x (must be 2^n for fast Fourier transform)
dx = Lx / Nx # Spatial step size
dt = 0.05 # Time step for recursive calculations
time_frames = np.arange(0.,600.,2.5)
Ep, x0, sigma = 0.1, -100, 15 # the initial particle's wavefubction parameters
barrier_pos, barrier_height, barrier_width = 0, 0.2, 3 # barrier parameters
# Make coordinate grid
x = np.linspace(-Lx / 2, Lx / 2, Nx) # x-space, make Nx points in the defined limits
# potential operator for half step (dt/2) in x-space
flatV = np.zeros(Nx) # make array with Nx points, set all elements to zero
flatV[int(Nx/2 + barrier_pos/dx) : int(Nx/2 + (barrier_pos+barrier_width)/dx)] = barrier_height
potential = np.exp(-1j * flatV / hbar * dt/2) # potential operator for half-step dt/2
# Kinetic energy operator for full step (dt) in momentum space
dk = 2 * np.pi / Lx # step in k-space grid corresponding to x-space
kx = dk * np.arange(-Nx/2, Nx/2) # k-sace grid
kx = fftshift(kx) # Shift for proper FFT
kinetic_operator = np.exp(-1j * hbar * kx**2 / (2 * m) * dt) # kinetic operator for full step
# Create the initial wavepacket (Gaussian)
psi0 = Gaussian_wave_at_E(x,x0 = x0,Ep = Ep,sigma=sigma)
# Create plot to animate
fig = plt.figure(figsize = (10,4) ,dpi=50) # make plot frame of defined size and dpi resolution (better set dpi to 100)
ax = plt.axes(xlim = (np.min(x), np.max(x)),ylim = (-.4,.3)) # define axes and x/y limits of graph
ax.plot(x,flatV,color = 'black') # plot potential
line_wfr, = ax.plot(x,psi0.real,color = 'red') # plot real part of wavefunction, keep reference to change it later
line_wfi, = ax.plot(x,psi0.imag,color = 'blue') # imaginary part of wavefunction
line_wf2, = ax.plot(x,np.square(np.abs(psi0))*5-.3,color = 'green') # suqare of the wavefunction
title = ax.set_title(f'Time: {0}') # define title, keep reference to change it in animation
ax.grid() # plot grid
ax.set_xlabel('X')
psi, curtime = psi0.copy(), 0 # reset initial wavefunction and current time when first start it
# Function to update the wavepacket for animation
def init():
global psi,curtime #define as global, so we can change these global parameters, nest time we run 'update' the previous psi and curtime will be used
psi, curtime = psi0.copy(), 0
def update(frame): #update animation, the FuncAnimation will call it to replot graph for each consequent 'frame'
print(f'\rFrame {frame+1} of {len(time_frames)} ',end = '')
global psi,curtime #define as global, so we can change these global parameters, nest time we run 'update' the previous psi and curtime will be used
if frame == 0: psi, curtime = psi0.copy(), 0 # reset initial wavefunction and current time when first start it
psi,curtime = GetPsi4(psi,curtime,time_frames[frame],potential,kinetic_operator,dt) #calculate new wavefunction for new time
line_wfr.set_data(x,psi.real) # replot these curves using updated psi
line_wfi.set_data(x,psi.imag)
line_wf2.set_data(x, np.square(np.abs(psi))*5-.4)
title.set_text(f'Time: {time_frames[frame]:.2f}')
return
# create animation object:
ani = FuncAnimation(fig, update, frames=len(time_frames), interval=30, init_func=init, blit=False) # setting blit=True does not replot map!
plt.close()
if save_mp4 == True:
name = f'packet_E{Ep}_barrier_h{barrier_height}-w{barrier_width}-TDSE-method2.mp4'
print(f'Saving in "{name}"...')
ani.save(name) # will save a movie in mp4 video file. Need to install ffmpeg codec and add its /bin/ to windows PATH!
print('done')
if show_html == True:
print('Making HTML...')
# ani.to_jshtml henerates animation in jshtml format; HTML incorporates the code into browser outptut window
# it will work only when you run it in html based python like Jupiter notebook
# you can use ani.save('movie.mp4') to sav it as mp4 file, provided that ffmpeg is installed separately with PATH to it.
#psi, curtime = psi0.copy(), 0 # reset initial wavefunction and current time before starting ani.jshtml
anim = ani.to_jshtml()
print('done')
display(HTML(anim))
Tunneling animation in 2D: potential barrier¶
This is animation of wavepacket propagation in two dimensions through a potential barrier.
def GetPsi3(psi0,t0,tf,potential,kinetic_operator,dt):
psi = psi0 # initial psi at tme t0. Note: this does not make a copy of psi0, but merely assigns name psi to the same array
while fabs(tf-t0) > dt/2 and t0 < tf: # continue recursively incrementing time by dt until time reaches tf or we are within dt/2 of target time
t0 += dt # new time point, increment by dt
# Apply potential operator (half step)
psi = np.exp(-1j * dt * potential / (2 * hbar)) * psi
# Apply kinetic operator (full step in Fourier space)
psi_k = np.fft.fft2(psi) # fourier transform into p=space (k-space)
psi_k *= kinetic_operator # propagate by dt in momentum space
psi = np.fft.ifft2(psi_k) # fourier inverse back to x-space
# Apply potential operator (half step)
psi = np.exp(-1j * dt * potential / (2 * hbar)) * psi
return psi, t0 # return wavefunction at new time and new time itself
# Simulation parameters
Lx, Ly = 100, 100 # Spatial dimensions of the coordinate space
Nx, Ny = 256, 256 # Number of grid points along x and y axes (must be 2^n for fast Fourier transform)
dx, dy = Lx / Nx, Ly / Ny # Spatial step sizes
dt = 0.05 # Time step for recursive calculations
time_frames = np.arange(0.,50.,.25)
barrier_width = 1
barrier_height = 1
# coordinate grid
x = np.linspace(-Lx / 2, Lx / 2, Nx)
y = np.linspace(-Ly / 2, Ly / 2, Ny)
X, Y = np.meshgrid(x, y)
# Initial 2D wavepacket (Gaussian)
sigma_x, sigma_y = 5, 5 # Width of the Gaussian
x0, y0 = -Lx/4, -Ly/4 # Initial position
px0, py0 = 1, 1 # 0.5 # Initial momentum
psi0 = np.exp(-((X - x0)**2 / (2 * sigma_x**2) + (Y - y0)**2 / (2 * sigma_y**2))) * \
np.exp(1j * (px0 * X + py0 * Y))
psi0 /= np.sqrt(np.sum(np.abs(psi0)**2) * dx * dy) # Normalize
psi = psi0.copy()
# Potential (e.g., a simple barrier in 2D)
V0 = barrier_height
potential = np.zeros((Nx,Ny))
#potential = V0 * (np.exp(-((X - 20)**2 + (Y - 0)**2) / 20) + np.exp(-((X + 20)**2 + (Y - 0)**2) / 20))
#potential = potential * 0
for j in range(int(Nx/2),int(Nx/2+barrier_width/dx)):
for i in range(len(potential)):
potential[i][j] = V0
# Wavenumber space
kx = 2 * np.pi * np.fft.fftfreq(Nx, d=dx) # these correspond to momentum eigenvalues
ky = 2 * np.pi * np.fft.fftfreq(Ny, d=dy)
Kx, Ky = np.meshgrid(kx, ky)
# Kinetic energy operator in momentum space
kinetic_operator = np.exp(-1j * dt * (hbar**2 / (2 * m)) * (Kx**2 + Ky**2))
fig, ax = plt.subplots(dpi = 50) # set higher resolution for smooth! dpi=200 recommended (default is 100)
imshow = ax.imshow(np.abs(psi)**2, extent=[-Lx/2, Lx/2, -Ly/2, Ly/2], origin='lower', cmap='viridis') # 2D color map of wavefunction
ax.contour(X, Y, potential, levels=[V0/2], colors='white', linestyles='solid',linewidths = 0.5)
title = ax.set_title(f'Time:')
ax.set_xlabel('X')
ax.set_ylabel('Y')
# Function to update the wavepacket for animation
curtime = 0
def init(): # initialize animation - get psi and curtime to 0
global psi,curtime # must be defined as global to change oputside variables
psi, curtime = psi0.copy(), 0
def update(frame):
print(f'\rframe={frame}/{len(time_frames)}: {time_frames[frame]}...',end = '')
global psi,curtime
psi,curtime = GetPsi3(psi,curtime,time_frames[frame],potential,kinetic_operator,dt)
imshow.set_data(np.abs(psi)**2)
imshow.set_clim(0,np.max(np.abs(psi)**2))
title.set_text(f'Time: {time_frames[frame]:.2f}')
return
ani = FuncAnimation(fig, update, frames=len(time_frames), init_func = init, interval=30, blit=False) # setting blit=True does not replot map!
plt.close()
if save_mp4 == True:
name = f'Tunneling2D_E{Ep}_V{V0}_w{barrier_width}-TDSE.mp4'
print(f'Saving in "{name}"...')
ani.save(name) # will save a movie in mp4 video file. Need to install ffmpeg codec and add its /bin/ to windows PATH!
print('done')
if show_html == True:
print('Making jshtml animation')
anim = ani.to_jshtml()
print('done')
display(HTML(anim))
Animation in 2D trap: harmonic and anharmonic potentials¶
This is animation of wavepacket dynamics in two dimensional harmonic and anharmonic potential. To pick specific potential set potential_type to 1,2 or 3 (or create your own).
It may take long time as currently there are 1000 time frames to calculate (see time_frames array). You may want to set same_mp4 to False if you only want movie in the browser.
potential_type = 2 # (1-harmonic, 2 very anharmonic, 3 add cubic term, 4 rough parabola)
#save_mp4 = False # (un)comment to override the True value to expedite calculations
from matplotlib import cm
# Simulation parameters
Lx, Ly = 100, 100 # Spatial dimensions of the coordinate space
Nx, Ny = 256, 256 # Number of grid points along x and y axes (must be 2^n for fast Fourier transform)
dx, dy = Lx / Nx, Ly / Ny # Spatial step sizes
dt = 0.05 # Time step for recursive calculations
time_frames = np.arange(0,500,1) # set max time to 200 for potential 1 (repeating), 1000 for potential 2
# coordinate grid
x = np.linspace(-Lx / 2, Lx / 2, Nx)
y = np.linspace(-Ly / 2, Ly / 2, Ny)
X, Y = np.meshgrid(x, y)
# Initial 2D wavepacket (Gaussian)
sigma_x, sigma_y = 10, 10 # Width of the Gaussian
x0, y0 = -15, 0 # Initial position
px0, py0 = 0, 1 # 0.5 # Initial momentum
psi = np.exp(-((X - x0)**2 / (2 * sigma_x**2) + (Y - y0)**2 / (2 * sigma_y**2))) * \
np.exp(1j * (px0 * X + py0 * Y))
psi /= np.sqrt(np.sum(np.abs(psi)**2) * dx * dy) # Normalize
psi0 = psi.copy()
Ep = (px0**2 + py0**2)/(2*m) # approx kinetic energy
# MAKE POTENTIALS : pick one by setting potential_type to 1,2,3
if potential_type == 1: # 1. Pure Harmonic potential
potential = (X**2 + Y**2)
pot = 'Parabolic V=x^2+y^2'
#potential = potential / potential[int(Lx/4)][0] * Ep * 4 # normalize to something...
elif potential_type == 2: # 2. strange anharmonicity
potential = (X**2 + Y**2) # boils down to A*(X^2 + Y^2) + (XY)^2 + Y^4
potential = potential / potential[int(Lx/4)]
pot = 'Strange V=(x^2+y^2) div (a + y^2)'
elif potential_type == 3: # 3. cubic anharmonicity along y
potential = (X**2 + Y**2 + 2 * Y**3/Ly + 5*Y**4/Ly**2)
pot = 'Anharmonic V=x^2+y^2+Ay^3+By^4'
elif potential_type == 4: # 4. Add roughness
potential = X**2 + Y**2 + 30*(np.sin(X/3) + np.sin(Y/3))**2 #30*np.sin(np.sqrt(X**2+Y**2)/2)**2 #10*(np.sin(X/2) + np.sin(Y/2))**2
pot = 'Rough V=x^2+y^2+A(sin(0.3x)+sin(0.3y))^2'
name = f"Trap2D_{pot}" # file name to save mp4 and png to
potential = potential / potential[int(Lx/4)][0] * Ep * 4 # normalize to something...
Epk = (px0**2 + py0**2)/(2*m) # approx kinetic energy (for precixe must do <psi|H|psi>, but I am lazy)
Epp = potential[int((Lx/2 + x0)/dx)][int((Ly/2 + y0)/dy)] # approx potential energy
print(f'Approx. particle energy T={Epk}, V={Epp}, E={Epk+Epp}')
#potential = V0 * (np.exp(-((X - 20)**2 + (Y - 0)**2) / 20) + np.exp(-((X + 20)**2 + (Y - 0)**2) / 20))
#potential = potential * 0
#plot potential in 3D
fig,ax = plt.subplots(dpi=50) # set dpi=200 for better resolution!
ax.set_axis_off()
ax = fig.add_subplot(111,projection='3d')
ax.plot_surface(X,Y,potential,cmap=cm.coolwarm, edgecolor = 'black',rcount=100,ccount=100,antialiased = True,lw=0.1)
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_title(pot)
plt.savefig(name+'.png')
plt.show()
plt.close()
# Wavenumber space
kx = 2 * np.pi * np.fft.fftfreq(Nx, d=dx) # these correspond to momentum eigenvalues
ky = 2 * np.pi * np.fft.fftfreq(Ny, d=dy)
Kx, Ky = np.meshgrid(kx, ky)
# Kinetic energy operator in momentum space
kinetic_operator = np.exp(-1j * dt * (hbar**2 / (2 * m)) * (Kx**2 + Ky**2))
plt.rcParams['animation.embed_limit'] = 500.0 # increase default max animation size limit to 500 MB
fig, ax = plt.subplots(dpi=50) # set dpi=200 for better resolution!
ax.set_aspect('equal')
imshow = ax.imshow(np.abs(psi)**2, extent=[-Lx/2, Lx/2, -Ly/2, Ly/2], origin='lower', cmap='viridis') # 2D color map of wavefunction
ax.contour(X, Y, potential, levels=[Epk+Epp], colors='white', linestyles='solid',linewidths = 0.5)
title = ax.set_title(f'Harmonic trap, Time: 0')
ax.set_xlabel('X')
ax.set_ylabel('Y')
# Function to update the wavepacket for animation
curtime = 0
def init():
global psi,curtime
psi, curtime = psi0.copy(), 0
def update(frame):
print(f'\rframe={frame+1} of {len(time_frames)}: t={time_frames[frame]}...',end = '')
global psi,curtime
psi,curtime = GetPsi3(psi,curtime,time_frames[frame],potential,kinetic_operator,dt)
imshow.set_data(np.abs(psi)**2)
imshow.set_clim(0,np.max(np.abs(psi)**2))
title.set_text(f'{pot} Time: {time_frames[frame]:.2f}')
return
ani = FuncAnimation(fig, update, init_func=init, frames=len(time_frames), interval=30, blit=False) # setting blit=True does not replot map!
plt.close()
if save_mp4 == True:
print(f'Making MP4 animation: {name}.mp4')
from matplotlib.animation import writers
Writer = writers['ffmpeg']
writer = Writer(metadata=dict(artist='Sergei Savikhin', genre = 'Quantum mechanics', title = 'Wavepacket in 2D '+pot),fps = 30)
ani.save(name+'.mp4', writer = writer)
print('done')
if show_html == True:
print('Making jshtml animation:')
anim = ani.to_jshtml()
print('done')
display(HTML(anim))
